{-# LANGUAGE ScopedTypeVariables,BangPatterns,DeriveDataTypeable #-}
module Random where

{- stolen from haskell.org - presumably copyright someone there ... -}
import Data.Array.ST
import Control.Monad
import Control.Monad.ST
import Data.STRef

import Data.List (sortBy)

import Statistics.RandomVariate as RV

import Data.Binary
import Data.Generics

import Control.DeepSeq


import Debug.Trace
mtrace o = (trace ("MTRACE : " ++ (show o) ++ "\n") o)




-- ProbabilityDistribution

data ProbabilityDistribution
    = ConstantPD Double |
      NormalPD Double Double | -- mean, sd
      LogNormalPD Double Double | -- ditto
      IntRoundPD ProbabilityDistribution | -- ditto
      UniformPD Double Double -- min, max
    deriving (Show,Data,Typeable)

instance Binary ProbabilityDistribution where
  put (ConstantPD a) = putWord8 0 >> put a
  put (NormalPD a b) = putWord8 1 >> put a >> put b
  put (LogNormalPD a b) = putWord8 2 >> put a >> put b
  put (IntRoundPD a) = putWord8 3 >> put a
  put (UniformPD a b) = putWord8 4 >> put a >> put b
  get = do
    tag_ <- getWord8
    case tag_ of
      0 -> get >>= \a -> return (ConstantPD a)
      1 -> get >>= \a -> get >>= \b -> return (NormalPD a b)
      2 -> get >>= \a -> get >>= \b -> return (LogNormalPD a b)
      3 -> get >>= \a -> return (IntRoundPD a)
      4 -> get >>= \a -> get >>= \b -> return (UniformPD a b)
      _ -> fail "no parse"


instance NFData ProbabilityDistribution where
  rnf (ConstantPD a) = rnf a
  rnf (NormalPD a b) = rnf a `seq` rnf b
  rnf (LogNormalPD a b) = rnf a `seq` rnf b
  rnf (IntRoundPD a) = rnf a
  rnf (UniformPD a b) = rnf a `seq` rnf b

-- Strictness since it's daft to leave tis hanging around
getDNum :: Gen a -> ProbabilityDistribution -> ST a Double
getDNum g (ConstantPD n) = return n
getDNum g (UniformPD lo hi) = 
    do (!d) <- RV.uniform g  -- saves a heap of heap
       return (((hi - lo) * d) + lo)
--getDNum g (NormalPD mean sd) = (mean,g')
--	where (_,g') = (randomR (0.0::Double,1.0) g)
-- Box - Muller transform 
--  should replace with native ziggurat at some stage 
-- 
getDNum g (NormalPD mean sd) 
        = do u <- RV.normal g
             let (!ret) = mean + u*sd
             return ret
getDNum g (IntRoundPD d) =
    do (!nd) <- getDNum g d
       return (fromIntegral $ round nd)

-- XXX - THIS NEXT BIT IS BROKEN AND FAILS SOME SIMPLE STATISTICAL TESTS!  
-- FOR |hi - lo| ~= 600 there's a spike in the histogram around 0
-- THIS IS FOR TESTING PURPOSES ONLY UNTIL I GET A DECENT IMPLEMENTATION 
randRange :: Gen a -> Int -> Int -> ST a Int
randRange g lo hi 
    | lo == hi = return lo
    | otherwise =
    do w::Int <- RV.uniform g          
       return (rn w)
    where rn :: Int -> Int
          rn n = (n `mod` ((hi - lo) + 1)) + lo   -- lo to hi inclusive


{-
-- Dodgy shuffle - for a speed test since proper shuffle is flogging us ..
-- Dodgy shuffle is in fact slower!!

shuffle :: Gen a -> [b] -> ST a [b]
shuffle g xs = 
    do rands::[Int] <- replicateM (length xs) (RV.uniform g)
       return $ map snd $ sortBy (\x y -> compare (fst x) (fst y)) $ zip rands xs

-}

shuffle :: Gen a -> [b] -> ST a [b]
shuffle g xs = do
        ar <- newArray n xs
        forM [1..n] $ \i -> do
            j <- RV.uniform g -- (i,n)
            vi <- readArray ar i
            vj <- readArray ar (rnum j i n)
            writeArray ar (rnum j i n) vi
            return vj
  where
    n = length xs
    newArray :: Int -> [a] -> ST s (STArray s Int a)
    newArray n xs =  newListArray (1,n) xs
    rnum num lo hi = (num `mod` (hi - lo + 1)) + lo -- modif for mersenne
-- XXX PREVIOUS LINE IS PROBABLY BROKEN!!!


{-
-- | Randomly shuffle a list without the IO Monad
--   /O(N)/
shuffle' :: [a] -> StdGen -> ([a],StdGen)
shuffle' xs gen = runST (do
        g <- newSTRef gen
        let randomRST lohi = do
              (a,s') <- liftM (randomR lohi) (readSTRef g)
              writeSTRef g s'
              return a
        ar <- newArray n xs
        xs' <- forM [1..n] $ \i -> do
                j <- randomRST (i,n)
                vi <- readArray ar i
                vj <- readArray ar j
                writeArray ar j vi
                return vj
        gen' <- readSTRef g
        return (xs',gen'))
  where
    n = length xs
    newArray :: Int -> [a] -> ST s (STArray s Int a)
    newArray n xs =  newListArray (1,n) xs
-}